from langgraph.graph import StateGraph, END
from src.agents.base import AnalystAgent, RiskAgent, AdvisorAgent
from src.core.llm_config import ModelRouter
import logging

# Initialize model router
try:
    router = ModelRouter()
    print("模型路由初始化成功")
except Exception as e:
    logging.error(f"模型路由初始化失败: {e}")
    raise

# Initialize agents with the model router
analyst_agent = AnalystAgent(router)
risk_agent = RiskAgent(router)
advisor_agent = AdvisorAgent(router)

# Define the state structure
class AgentState:
    def __init__(self, user_input, selected_role="全流程"):
        self.user_input = user_input
        self.selected_role = selected_role
        self.analysis_result = None
        self.risk_assessment = None
        self.final_advice = None

# Define node functions with error handling
def analyst_node(state):
    try:
        print(f"分析师Agent处理: {state.user_input}")
        state.analysis_result = analyst_agent.analyze_market_trends(state.user_input)
        return state
    except Exception as e:
        logging.error(f"分析师节点处理失败: {e}")
        state.analysis_result = f"分析失败: {str(e)}"
        return state

def risk_node(state):
    try:
        print(f"风控师Agent处理: {state.analysis_result}")
        state.risk_assessment = risk_agent.assess_risk({
            'query': state.user_input,
            'analysis': state.analysis_result
        })
        return state
    except Exception as e:
        logging.error(f"风控节点处理失败: {e}")
        state.risk_assessment = f"风险评估失败: {str(e)}"
        return state

def advisor_node(state):
    try:
        print(f"顾问Agent处理: {state.risk_assessment}")
        state.final_advice = advisor_agent.provide_advice({
            'query': state.user_input,
            'assessment': state.risk_assessment
        })
        return state
    except Exception as e:
        logging.error(f"顾问节点处理失败: {e}")
        state.final_advice = f"建议生成失败: {str(e)}"
        return state

# Create workflow graph
workflow = StateGraph(AgentState)

# Add nodes
workflow.add_node("analyst", analyst_node)
workflow.add_node("risk", risk_node)
workflow.add_node("advisor", advisor_node)

# Define edges
workflow.set_entry_point("analyst")
workflow.add_edge("analyst", "risk")
workflow.add_edge("risk", "advisor")
workflow.add_edge("advisor", END)

# Compile the graph
agent_workflow = workflow.compile()

# Example usage
if __name__ == "__main__":
    input_query = "AAPL股票投资分析"
    try:
        result = agent_workflow.invoke(AgentState(input_query))
        print("\n最终投资建议:")
        print(result.final_advice)
    except Exception as e:
        logging.error(f"工作流执行失败: {e}")
        print("工作流执行出错，请检查日志")